import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import Runnable, RunnablePassthrough
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import AIMessage, HumanMessage, BaseMessage
from typing import List
from pydantic import ConfigDict

# --- 环境准备 ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码 (修复版) ---

    llm = ChatOpenAI(model=model_name)
    print(f"模型已初始化 (使用 {model_name})。")

    # 1. 创建"纪要员"链
    # 它的任务是把对话历史总结成一句话
    summarizer_prompt = ChatPromptTemplate.from_messages([
        ("user", "请将以下对话内容提炼成一句简洁的摘要，以保留核心信息：\n{chat_history}"),
    ])
    summarizer = summarizer_prompt | llm | StrOutputParser()
    print("纪要员链 (Summarizer) 已创建。")

    # 2. 创建一个修正后的、带摘要功能的历史记录类
    # 继承自 ChatMessageHistory，这是一个 Pydantic 模型
    class SummarizedChatMessageHistory(ChatMessageHistory):
        # 允许 Pydantic 模型包含像 Runnable 这样的复杂/任意类型
        model_config = ConfigDict(arbitrary_types_allowed=True)
        
        # 正确地将 summarizer 和 max_messages 声明为类的字段
        summarizer: Runnable
        max_messages: int

        def add_messages(self, messages: List[BaseMessage]) -> None:
            """在添加新消息后，检查是否需要触发摘要。"""
            super().add_messages(messages)
            
            if len(self.messages) > self.max_messages:
                print("\n--- 触发摘要 ---")
                # 将要被压缩成摘要的旧消息
                messages_to_summarize = self.messages[:-(self.max_messages//2)]
                # 调用"纪要员"链来生成摘要
                summary = self.summarizer.invoke({"chat_history": messages_to_summarize})
                print(f"生成的新摘要: {summary}")
                # 创建一条新的消息来存放摘要
                summary_message = AIMessage(content=f"这是之前的对话摘要: {summary}")
                # 保留最近的几条消息
                latest_messages = self.messages[-(self.max_messages//2):]
                # 用"摘要 + 最新消息"来更新整个历史记录
                self.messages[:] = [summary_message] + latest_messages
                print("--- 摘要完成 ---")

    # 3. 创建主对话链的提示
    prompt = ChatPromptTemplate.from_messages([
        # 注意：我们不再直接用 "system" 消息，因为摘要本身会作为一条 AIMessage 出现
        MessagesPlaceholder(variable_name="chat_history"),
        ("user", "{input}"),
    ])
    
    # 4. 重新引入 store 和 get_session_history 函数来正确管理会话
    store = {}
    def get_session_history(session_id: str) -> SummarizedChatMessageHistory:
        if session_id not in store:
            # 如果是新会话，创建一个我们自定义的、带摘要功能的历史对象
            print(f"为新会话 '{session_id}' 创建记忆体...")
            store[session_id] = SummarizedChatMessageHistory(
                summarizer=summarizer, 
                max_messages=4 # 为了演示，设一个小一点的窗口
            )
        return store[session_id]

    # 5. 用 RunnableWithMessageHistory 包装，并使用我们的 get_session_history 函数
    with_message_history = RunnableWithMessageHistory(
        prompt | llm | StrOutputParser(),
        get_session_history, # 使用正确的会话管理函数
        input_messages_key="input",
        history_messages_key="chat_history",
    )
    print("带自定义摘要记忆的链已创建！")

    # 6. 开始对话
    session_id = "chat_with_lcel_summary"
    print("\n--- 开始对话 (将在超过4条消息后触发摘要) ---")
    
    inputs = [
        "你好，我叫派大星，我的朋友是海绵宝宝。",
        "我们住在比奇堡的水下菠萝屋里。",
        "我最喜欢的活动是抓水母。",
        "海绵宝宝最喜欢做什么？", # 此时应该会触发摘要
        "你还记得我叫什么吗？" # 验证摘要是否保留了信息
    ]

    for i, user_input in enumerate(inputs):
        print(f"\n第 {i+1} 轮\n[你]: {user_input}")
        response = with_message_history.invoke(
            {"input": user_input},
            config={"configurable": {"session_id": session_id}}
        )
        print(f"[AI]: {response}")

    print("\n--- 对话结束 ---") 